import os
import json
from glob import glob


webqsp_train = []
for p in glob("/home/xionggm/codes/LLM_KGQA/save-anno-clean/webqsp/**/*.json"):
    webqsp_train.append(json.load(open(p)))

cwq_train = []
for p in glob("/home/xionggm/codes/LLM_KGQA/save-anno-clean/cwq/**/*.json"):
    cwq_train.append(json.load(open(p)))

webqsp_train_ids = [d["id"].replace(".P0","") for d in webqsp_train]
cwq_train_ids = [d["id"].replace(".P0","") for d in cwq_train]

webqsp_train_ids = set(webqsp_train_ids)
cwq_train_ids = set(cwq_train_ids)

webqsp_test = []
for p in glob("/home/xionggm/codes/LLM_KGQA/data/webqsp/test/*.json"):
    with open(p) as f:
        webqsp_test += json.load(f)

cwq_test = []
for p in glob("/home/xionggm/codes/LLM_KGQA/data/cwq/test/*.json"):
    with open(p) as f:
        cwq_test += json.load(f)

webqsp_test_ids = [d["id"].replace(".P0","") for d in webqsp_test]
cwq_test_ids = [d["id"].replace(".P0","") for d in cwq_test]

paths = [
    "save_dir/Retrieval/pyserini/search_results/QA_WebQSP_Freebase_BM25/train_fid_SPQA.json",
    "save_dir/Retrieval/pyserini/search_results/QA_WebQSP_Freebase_BM25/test_fid_SPQA.json",
    "save_dir/Retrieval/pyserini/search_results/QA_CWQ_Freebase_BM25/train_fid_SPQA.json",
    "save_dir/Retrieval/pyserini/search_results/QA_CWQ_Freebase_BM25/dev_fid_SPQA.json",
    "save_dir/Retrieval/pyserini/search_results/QA_CWQ_Freebase_BM25/test_fid_SPQA.json",
]
all_data = []
for p in paths:
    data = json.load(open(p))
    all_data += data

webqsp_train = [d for d in all_data if d["id"].replace(":QA","").replace(":SP","") in webqsp_train_ids]
cwq_train = [d for d in all_data if d["id"].replace(":QA","").replace(":SP","") in cwq_train_ids]
webqsp_test = [d for d in all_data if d["id"].replace(":QA","").replace(":SP","") in webqsp_test_ids]
cwq_test = [d for d in all_data if d["id"].replace(":QA","").replace(":SP","") in cwq_test_ids]

# merge
merge_train = webqsp_train + cwq_train
merge_test = webqsp_test + cwq_test

# save
outdir = "save_dir_fid_smalldata"
os.makedirs(outdir, exist_ok=True)
json.dump(webqsp_train, open(f"{outdir}/webqsp_train_SPQA.json", "w"), indent=4, ensure_ascii=False)
json.dump(cwq_train, open(f"{outdir}/cwq_train_SPQA.json", "w"), indent=4, ensure_ascii=False)
json.dump(webqsp_test, open(f"{outdir}/webqsp_test_SPQA.json", "w"), indent=4, ensure_ascii=False)
json.dump(cwq_test, open(f"{outdir}/cwq_test_SPQA.json", "w"), indent=4, ensure_ascii=False)
json.dump(merge_train, open(f"{outdir}/merge_train_SPQA.json", "w"), indent=4, ensure_ascii=False)
json.dump(merge_test, open(f"{outdir}/merge_test_SPQA.json", "w"), indent=4, ensure_ascii=False)


